package com.hzh.MLlib

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
import org.apache.spark.sql.{DataFrame, SparkSession}

object DemoTrainImageML {
  def main(args: Array[String]): Unit = {

    /**
     * 创建环境
     *
     */

    val spark: SparkSession = SparkSession
      .builder()
      .config("spark.sql.shuffle.partitions", 1)
      .master("local[6]")
      .appName("DemoTrainImageML")
      .getOrCreate()

    import org.apache.spark.sql.functions._
    import spark.implicits._
    /**
     * 1、读取数据
     *
     */
    val imageDataDF: DataFrame = spark
      .read
      .format("libsvm")
      .option("numFeatures", 784) //指定特征的长度
      .load("data/image_data")


    /**
     * 2、将数据集拆分成训练集和测试集
     *
     */
    val Array(train:DataFrame,test:DataFrame) = imageDataDF.randomSplit(Array(0.75, 0.25))

    /**
     *
     * 3、选择算法
     */

    //多项逻辑回归
    val logisticRegression = new LogisticRegression()

    /**
     * 4、将训练集带入算法训练模型
     * 
     */
    val logisticRegressionModel: LogisticRegressionModel = logisticRegression.fit(train)

    /**
     * 5、评估模型准确率
     *
     */
    val testDF: DataFrame = logisticRegressionModel.transform(test)

    testDF.cache()

    val p: Double = testDF.where($"label" === $"prediction").count().toDouble / testDF.count()


    println(s"模型的准确率：$p")

    logisticRegressionModel
      .write
      .overwrite()
      .save("data/image_model")
  }
}
